import os
import requests
import json
import base64
from io import BytesIO
from PIL import Image
import logging
from typing import Optional

logger = logging.getLogger(__name__)

class AIImageGenerator:
    def __init__(self, huggingface_token: Optional[str] = None):
        self.huggingface_token = huggingface_token or os.getenv("HUGGINGFACE_API_KEY")
        self.base_url = "https://api-inference.huggingface.co/models/"
        
    def generate_image(self, prompt: str, model: str = "stabilityai/stable-diffusion-2-1") -> Optional[Image.Image]:
        """使用HuggingFace API生成图像"""
        if not self.huggingface_token:
            logger.warning("HuggingFace API密钥未设置，无法生成AI图像")
            return None
            
        try:
            headers = {"Authorization": f"Bearer {self.huggingface_token}"}
            
            # 准备请求数据
            payload = {
                "inputs": prompt,
                "parameters": {
                    "num_inference_steps": 20,
                    "guidance_scale": 7.5
                }
            }
            
            logger.info(f"🖼️ 使用模型 {model} 生成图像: {prompt}")
            
            # 发送请求
            response = requests.post(
                f"{self.base_url}{model}",
                headers=headers,
                json=payload,
                timeout=60
            )
            
            if response.status_code == 200:
                # 将响应内容转换为图像
                image = Image.open(BytesIO(response.content))
                logger.info(f"✅ AI图像生成成功: {image.size}")
                return image
            else:
                logger.error(f"❌ AI图像生成失败: {response.status_code} - {response.text}")
                return None
                
        except Exception as e:
            logger.error(f"❌ AI图像生成异常: {str(e)}")
            return None
    
    def generate_image_for_scene(self, scene_description: str, style: str = "realistic") -> Optional[Image.Image]:
        """为场景生成合适的图像"""
        # 根据场景描述和风格优化提示
        enhanced_prompt = self._enhance_prompt(scene_description, style)
        
        # 选择适合的模型
        model = self._select_model(style)
        
        return self.generate_image(enhanced_prompt, model)
    
    def _enhance_prompt(self, description: str, style: str) -> str:
        """增强提示词以获得更好的生成效果"""
        style_keywords = {
            "realistic": "photorealistic, high quality, detailed, 8k",
            "artistic": "artistic, painting style, creative, vivid colors",
            "anime": "anime style, Japanese animation, vibrant colors",
            "digital_art": "digital art, concept art, fantasy, detailed",
            "cinematic": "cinematic, movie still, dramatic lighting, film quality"
        }
        
        style_text = style_keywords.get(style, "high quality, detailed")
        
        # 限制提示长度
        if len(description) > 500:
            description = description[:500]
            
        return f"{description}, {style_text}"
    
    def _select_model(self, style: str) -> str:
        """根据风格选择合适的模型"""
        model_mapping = {
            "realistic": "stabilityai/stable-diffusion-2-1",
            "artistic": "runwayml/stable-diffusion-v1-5",
            "anime": "22h/vintedois-diffusion-v0-1",
            "digital_art": "prompthero/openjourney",
            "cinematic": "stabilityai/stable-diffusion-2-1"
        }
        
        return model_mapping.get(style, "stabilityai/stable-diffusion-2-1")